"""
OCR智能体
处理图像OCR识别和文本提取的智能体
"""

from typing import Dict, Any, Optional, List
import asyncio
import base64
import io
from PIL import Image
import pytesseract

from apps.agents.base_agent import BaseAgent, AgentType
from utils.logger import logger


class OCRAgent(BaseAgent):
    """OCR识别智能体"""
    
    def __init__(self, agent_id: str = "ocr_agent"):
        super().__init__(
            agent_id=agent_id,
            agent_type=AgentType.OCR,
            name="OCR识别智能体"
        )
        
        # OCR配置
        self.supported_formats = ["jpg", "jpeg", "png", "bmp", "tiff", "pdf"]
        self.max_file_size = 10 * 1024 * 1024  # 10MB
        self.default_language = "chi_sim+eng"  # 中英文
        
        logger.info(f"OCR智能体初始化完成: {self.agent_id}")
    
    async def validate_input(self, input_data: Dict[str, Any]) -> bool:
        """验证输入数据"""
        # 检查必需字段
        if "image_data" not in input_data and "image_path" not in input_data:
            logger.error("缺少图像数据或路径")
            return False
        
        # 验证图像数据
        if "image_data" in input_data:
            image_data = input_data["image_data"]
            if not isinstance(image_data, (str, bytes)):
                logger.error("图像数据格式无效")
                return False
            
            # 检查数据大小
            if isinstance(image_data, str):
                # base64字符串
                try:
                    decoded_data = base64.b64decode(image_data)
                    if len(decoded_data) > self.max_file_size:
                        logger.error("图像文件过大")
                        return False
                except Exception as e:
                    logger.error(f"base64解码失败: {str(e)}")
                    return False
        
        # 验证文件路径
        if "image_path" in input_data:
            image_path = input_data["image_path"]
            if not isinstance(image_path, str):
                logger.error("图像路径格式无效")
                return False
            
            # 检查文件格式
            file_ext = image_path.lower().split('.')[-1]
            if file_ext not in self.supported_formats:
                logger.error(f"不支持的文件格式: {file_ext}")
                return False
        
        return True
    
    async def process_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理OCR识别任务"""
        
        logger.info(f"开始处理OCR任务")
        
        try:
            # 1. 获取图像数据
            image = await self._load_image(task_data)
            
            # 2. 预处理图像
            processed_image = await self._preprocess_image(image, task_data)
            
            # 3. 执行OCR识别
            ocr_result = await self._perform_ocr(processed_image, task_data)
            
            # 4. 后处理结果
            final_result = await self._postprocess_result(ocr_result, task_data)
            
            logger.info(f"OCR任务完成，识别到 {len(final_result.get('text', ''))} 个字符")
            
            return final_result
            
        except Exception as e:
            logger.error(f"OCR任务失败: {str(e)}")
            raise
    
    async def _load_image(self, task_data: Dict[str, Any]) -> Image.Image:
        """加载图像"""
        if "image_data" in task_data:
            # 从base64数据加载
            image_data = task_data["image_data"]
            if isinstance(image_data, str):
                image_bytes = base64.b64decode(image_data)
            else:
                image_bytes = image_data
            
            return Image.open(io.BytesIO(image_bytes))
        
        elif "image_path" in task_data:
            # 从文件路径加载
            image_path = task_data["image_path"]
            return Image.open(image_path)
        
        else:
            raise ValueError("没有提供图像数据")
    
    async def _preprocess_image(self, image: Image.Image, task_data: Dict[str, Any]) -> Image.Image:
        """预处理图像"""
        
        # 转换为RGB模式（如果不是的话）
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # 调整图像大小（如果太大）
        max_dimension = task_data.get("max_dimension", 2048)
        if max(image.size) > max_dimension:
            ratio = max_dimension / max(image.size)
            new_size = tuple(int(dim * ratio) for dim in image.size)
            image = image.resize(new_size, Image.Resampling.LANCZOS)
        
        # 应用图像增强（如果启用）
        if task_data.get("enhance_image", False):
            image = await self._enhance_image(image)
        
        return image
    
    async def _enhance_image(self, image: Image.Image) -> Image.Image:
        """增强图像质量"""
        # 这里可以实现更复杂的图像增强算法
        # 目前只做简单的对比度增强
        from PIL import ImageEnhance
        
        # 增强对比度
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(1.2)
        
        # 增强锐度
        enhancer = ImageEnhance.Sharpness(image)
        image = enhancer.enhance(1.1)
        
        return image
    
    async def _perform_ocr(self, image: Image.Image, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """执行OCR识别"""
        
        # 获取OCR参数
        language = task_data.get("language", self.default_language)
        oem = task_data.get("oem", 3)  # OCR Engine Mode
        psm = task_data.get("psm", 6)  # Page Segmentation Mode
        
        # OCR配置
        custom_config = f'--oem {oem} --psm {psm} -l {language}'
        
        try:
            # 执行OCR
            text = pytesseract.image_to_string(image, config=custom_config)
            
            # 获取详细信息（如果启用）
            detailed_info = None
            if task_data.get("detailed_info", False):
                detailed_info = pytesseract.image_to_data(image, config=custom_config, output_type=pytesseract.Output.DICT)
            
            return {
                "text": text.strip(),
                "language": language,
                "detailed_info": detailed_info,
                "confidence": self._estimate_confidence(text, detailed_info)
            }
            
        except Exception as e:
            logger.error(f"OCR识别失败: {str(e)}")
            raise RuntimeError(f"OCR识别失败: {str(e)}")
    
    def _estimate_confidence(self, text: str, detailed_info: Optional[Dict] = None) -> float:
        """估计识别置信度"""
        
        if not text:
            return 0.0
        
        # 基础置信度计算
        confidence = 0.7  # 基础置信度
        
        # 根据文本长度调整
        if len(text) < 10:
            confidence *= 0.8  # 短文本置信度较低
        elif len(text) > 100:
            confidence *= 1.1  # 长文本置信度较高
        
        # 如果有详细信息，使用平均置信度
        if detailed_info and "conf" in detailed_info:
            confidences = [conf for conf in detailed_info["conf"] if conf != -1]
            if confidences:
                avg_confidence = sum(confidences) / len(confidences)
                confidence = max(confidence, avg_confidence / 100.0)
        
        # 检查文本质量
        if self._check_text_quality(text):
            confidence *= 1.1
        else:
            confidence *= 0.9
        
        return min(confidence, 1.0)
    
    def _check_text_quality(self, text: str) -> bool:
        """检查文本质量"""
        
        # 检查是否有过多的特殊字符
        special_chars = sum(1 for char in text if not char.isalnum() and not char.isspace())
        if special_chars / len(text) > 0.3:
            return False
        
        # 检查是否有合理的字符分布
        if len(text) > 0:
            chinese_ratio = sum(1 for char in text if '\u4e00' <= char <= '\u9fff') / len(text)
            english_ratio = sum(1 for char in text if char.isalpha() and char.isascii()) / len(text)
            
            # 如果中文和英文比例都很低，可能是识别错误
            if chinese_ratio < 0.1 and english_ratio < 0.1:
                return False
        
        return True
    
    async def _postprocess_result(self, ocr_result: Dict[str, Any], task_data: Dict[str, Any]) -> Dict[str, Any]:
        """后处理OCR结果"""
        
        text = ocr_result["text"]
        
        # 文本清理
        if task_data.get("clean_text", True):
            text = self._clean_text(text)
        
        # 格式化输出
        output_format = task_data.get("output_format", "text")
        
        if output_format == "json":
            return {
                "text": text,
                "confidence": ocr_result["confidence"],
                "language": ocr_result["language"],
                "word_count": len(text.split()),
                "char_count": len(text),
                "detailed_info": ocr_result.get("detailed_info")
            }
        
        else:  # text format
            return {
                "text": text,
                "confidence": ocr_result["confidence"],
                "language": ocr_result["language"]
            }
    
    def _clean_text(self, text: str) -> str:
        """清理识别文本"""
        
        # 移除多余的空白字符
        lines = text.split('\n')
        cleaned_lines = []
        
        for line in lines:
            line = line.strip()
            if line:  # 只保留非空行
                cleaned_lines.append(line)
        
        # 重新组合
        cleaned_text = '\n'.join(cleaned_lines)
        
        # 移除常见的OCR错误
        cleaned_text = cleaned_text.replace('|', 'I')  # 竖线常被识别为I
        cleaned_text = cleaned_text.replace('‘', "'")  # 中文引号
        cleaned_text = cleaned_text.replace('’', "'")
        cleaned_text = cleaned_text.replace('“', '"')
        cleaned_text = cleaned_text.replace('”', '"')
        
        return cleaned_text
    
    async def process_batch_ocr(self, image_list: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """批量处理OCR任务"""
        
        logger.info(f"开始批量OCR处理，共 {len(image_list)} 张图像")
        
        results = []
        
        for i, image_data in enumerate(image_list):
            try:
                logger.info(f"处理第 {i+1} 张图像")
                
                # 创建单个任务
                task_data = {
                    **image_data,
                    "batch_index": i
                }
                
                result = await self.process_task(task_data)
                result["batch_index"] = i
                results.append(result)
                
            except Exception as e:
                logger.error(f"第 {i+1} 张图像处理失败: {str(e)}")
                results.append({
                    "batch_index": i,
                    "error": str(e),
                    "text": "",
                    "confidence": 0.0
                })
        
        logger.info(f"批量OCR处理完成")
        return results